#define WIN32_LEAN_AND_MEAN

#include <winsock2.h>
#include <ws2tcpip.h>
#include <time.h>
#include <stdlib.h>
#include <stdio.h>
#include <direct.h>

#include "rdistro.h"
#include "../../rdx/src/lua/src/lua.h"
#include "../tdp2/tdp2_hash.hpp"

using namespace RDistro;

struct ADDRINFO_Cluster
{
	ADDRINFO *first;
};

int sha256(lua_State *L)
{
	size_t len;
	const char *str = lua_tolstring(L, 1, &len);
	UInt8 output[32];
	
	{
		TDP::Crypto::SHA256::SHA256State shaState;
		shaState.FeedBytes(str, len);
		shaState.Terminate();
		shaState.Flush(output);
	}

	lua_pushlstring(L, reinterpret_cast<const char *>(output), 32);
	return 1;
}

int Sockets_EnvironmentStartup(lua_State *L)
{
	srand(static_cast<unsigned int>(time(NULL)));

	WSADATA wsaData;

	int wsaStartupCode = WSAStartup(MAKEWORD(2, 2), &wsaData);
	if(wsaStartupCode != 0)
	{
		WSACleanup();
		lua_pushstring(L, "WS startup failed");
		lua_error(L);
		return 0;
	}

	lua_newuserdata(L, 1);

	return 1;
}

int Sockets_EnvironmentCleanup(lua_State *L)
{
	WSACleanup();
	return 0;
}

// SocketType, Server, Port, Binding
int Sockets_GetAddrInfoCluster(lua_State *L)
{
	ADDRINFO hints;
	ADDRINFO *addrInfo;
	memset(&hints, 0, sizeof(hints));
	hints.ai_family = PF_UNSPEC;
	hints.ai_socktype = SOCK_STREAM;

	if(lua_toboolean(L, 4) != 0)
		hints.ai_flags |= (AI_PASSIVE | AI_NUMERICHOST);

	const char *sockType = lua_tostring(L, 1);

	if(sockType != NULL && !strcmp(sockType, "UDP"))
		hints.ai_socktype = SOCK_DGRAM;
	int rv;
	rv = getaddrinfo(lua_tostring(L, 2), lua_tostring(L, 3), &hints, &addrInfo);

	if(rv != 0)
		return 0;
	
	ADDRINFO_Cluster *cluster = static_cast<ADDRINFO_Cluster *>(lua_newuserdata(L, sizeof(ADDRINFO_Cluster)));
	cluster->first = addrInfo;

	return 1;
}

int Sockets_AddrInfoClusterFirst(lua_State *L)
{
	ADDRINFO_Cluster *cluster = static_cast<ADDRINFO_Cluster *>(lua_touserdata(L, 1));
	if(cluster->first)
		lua_pushlightuserdata(L, cluster->first);
	else
		lua_pushnil(L);
	return 1;
}

int Sockets_AddrInfoNext(lua_State *L)
{
	ADDRINFO *ai = static_cast<ADDRINFO *>(lua_touserdata(L, 1));
	if(ai->ai_next)
		lua_pushlightuserdata(L, ai->ai_next);
	else
		lua_pushnil(L);
	return 1;
}

int Sockets_AddrInfoFamily(lua_State *L)
{
	ADDRINFO *ai = static_cast<ADDRINFO *>(lua_touserdata(L, 1));
	lua_pushinteger(L, ai->ai_family);
	return 1;
}

int Sockets_FreeAddrInfoCluster(lua_State *L)
{
	freeaddrinfo(static_cast<ADDRINFO_Cluster *>(lua_touserdata(L, 1))->first);
	return 0;
}

int Sockets_CreateSocket(lua_State *L)
{
	ADDRINFO *addrInfo = static_cast<ADDRINFO *>(lua_touserdata(L, 1));
	SOCKET s = socket(addrInfo->ai_family, addrInfo->ai_socktype, addrInfo->ai_protocol);
	if(s == INVALID_SOCKET)
		return 0;

	u_long mode = 1;
	ioctlsocket(s, FIONBIO, &mode);

	SOCKET *ps = static_cast<SOCKET *>(lua_newuserdata(L, sizeof(SOCKET)));
	*ps = s;
	return 1;
}

int Sockets_CloseSocket(lua_State *L)
{
	SOCKET *ps = static_cast<SOCKET *>(lua_touserdata(L, 1));
	closesocket(*ps);
	return 0;
}

int Sockets_Bind(lua_State *L)
{
	SOCKET *ps = static_cast<SOCKET *>(lua_touserdata(L, 1));
	ADDRINFO *ai = static_cast<ADDRINFO *>(lua_touserdata(L, 2));
	int bindResult = bind(*ps, ai->ai_addr, ai->ai_addrlen);
	if(bindResult == SOCKET_ERROR)
	{
		printf("Error attempting to bind socket! %i\n", WSAGetLastError());
		return 0;
	}
	lua_pushboolean(L, 1);
	return 1;
}

int Sockets_Listen(lua_State *L)
{
	SOCKET *ps = static_cast<SOCKET *>(lua_touserdata(L, 1));
	int backlog = static_cast<int>(lua_tointeger(L, 2));
	int result = listen(*ps, backlog);
	if(result == SOCKET_ERROR)
		return 0;
	lua_pushboolean(L, 1);
	return 1;
}

int Sockets_Shutdown(lua_State *L)
{
	SOCKET *ps = static_cast<SOCKET *>(lua_touserdata(L, 1));
	shutdown(*ps, SD_SEND);
	return 0;
}

// Socket
// Returns new socket, sockaddr
int Sockets_Accept(lua_State *L)
{
	struct sockaddr addr;
	int szAddr = sizeof(addr);
	SOCKET *ps = static_cast<SOCKET *>(lua_touserdata(L, 1));
	SOCKET newS = accept(*ps, &addr, &szAddr);

	if(newS == INVALID_SOCKET)
	{
		return 0;
	}

	SOCKET *pnews = static_cast<SOCKET *>(lua_newuserdata(L, sizeof(SOCKET)));
	*pnews = newS;
	struct sockaddr *paddr = static_cast<struct sockaddr *>(lua_newuserdata(L, sizeof(struct sockaddr)));
	*paddr = addr;
	return 2;
}

int Sockets_Send(lua_State *L)
{
	SOCKET *ps = static_cast<SOCKET *>(lua_touserdata(L, 1));
	const char *bytes = lua_tostring(L, 2);

	int rv = send(*ps, bytes, lua_objlen(L, 2), 0);
	return 0;
}

// Socket, length
int Sockets_TryReceive(lua_State *L)
{
	SOCKET *ps = static_cast<SOCKET *>(lua_touserdata(L, 1));
	char buf[70000];
	SOCKET s = *ps;

	int l = static_cast<int>(lua_tointeger(L, 2));

	if(l < 1 || l > sizeof(buf))
		return 0;

	int rv = recv(s, buf, l, 0);

	if(rv == SOCKET_ERROR && WSAGetLastError() == WSAEWOULDBLOCK)
	{
		lua_pushboolean(L, 1);
		lua_pushnil(L);
		return 2;
	}

	if(rv == SOCKET_ERROR)
	{
		lua_pushboolean(L, 0);
		lua_pushinteger(L, WSAGetLastError());
		return 2;
	}
	
	lua_pushboolean(L, 1);
	lua_pushlstring(L, buf, rv);
	return 2;
}

int Sockets_SelectSingle(lua_State *L)
{
	SOCKET *ps = static_cast<SOCKET *>(lua_touserdata(L, 1));
	lua_Integer grouping = lua_tointeger(L, 2);

	TIMEVAL timeout;
	timeout.tv_usec = 0;
	timeout.tv_sec = 0;

	fd_set sockSet;
	sockSet.fd_array[0] = *ps;
	sockSet.fd_count = 1;

	fd_set *sets[3];
	sets[0] = sets[1] = sets[2] = NULL;

	sets[grouping] = &sockSet;

	int nSelected = select(1, sets[0], sets[1], sets[2], &timeout);

	lua_pushboolean(L, (nSelected > 0) ? 1 : 0);
	return 1;
}

int ForceSetMetaTable(lua_State *L)
{
	lua_pushvalue(L, 2);
	lua_setmetatable(L, 1);
	return 0;
}

int SleepThread(lua_State *L)
{
	Sleep(static_cast<DWORD>(lua_tointeger(L, 1)));
	return 0;
}

int EncodeInteger(lua_State *L)
{
	UInt8 bytes[32];
	lua_Integer li = lua_tointeger(L, 1);
	size_t numBytes = static_cast<size_t>(lua_tointeger(L, 2));
	for(int i=0;i<numBytes;i++)
		bytes[i] = static_cast<UInt8>((li >> (i*8)) & 0xff);
	lua_pushlstring(L, reinterpret_cast<const char *>(bytes), numBytes);
	return 1;
}

int DecodeInteger(lua_State *L)
{
	const char *str = lua_tostring(L, 1);
	size_t numBytes = lua_objlen(L, 1);

	lua_Integer li = 0;
	for(int i=0;i<numBytes;i++)
		li |= ((static_cast<lua_Integer>(str[i]) & 0xFF) << (i*8));
	lua_pushinteger(L, li);
	return 1;
}

int TimeSeconds(lua_State *L)
{
	lua_pushinteger(L, static_cast<lua_Integer>(clock() / CLOCKS_PER_SEC));
	return 1;
}

int GenerateContentKey(lua_State *L)
{
	const char *contentChars = "0123456789abcdefghijklmnopqrstuvwxyz";
	size_t possibleChars = strlen(contentChars);
	char key[32];
	for(int i=0;i<32;i++)
		key[i] = contentChars[rand() % possibleChars];
	lua_pushlstring(L, key, 32);
	return 1;
}

int GenerateAuthKey(lua_State *L)
{
	char baseKey[32];
	for(int i=0;i<32;i++)
		baseKey[i] = (rand() % 256);
	UInt8 hashedKey[32];

	{
		TDP::Crypto::SHA256::SHA256State shaState;
		shaState.FeedBytes(baseKey, sizeof(baseKey));
		shaState.Terminate();
		shaState.Flush(hashedKey);
	}
	lua_pushlstring(L, reinterpret_cast<const char *>(hashedKey), 32);
	return 1;
}

int FileTimestamp(lua_State *L)
{
	UInt8 timestamp[8];

	HANDLE f = CreateFile(lua_tostring(L, 1), GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
	if(f == INVALID_HANDLE_VALUE)
		return 0;

	FILETIME modTime;
	GetFileTime(f, NULL, NULL, &modTime);
	CloseHandle(f);

	timestamp[0] = static_cast<UInt8>(modTime.dwHighDateTime >> 24);
	timestamp[1] = static_cast<UInt8>(modTime.dwHighDateTime >> 16);
	timestamp[2] = static_cast<UInt8>(modTime.dwHighDateTime >> 8);
	timestamp[3] = static_cast<UInt8>(modTime.dwHighDateTime);
	timestamp[4] = static_cast<UInt8>(modTime.dwLowDateTime >> 24);
	timestamp[5] = static_cast<UInt8>(modTime.dwLowDateTime >> 16);
	timestamp[6] = static_cast<UInt8>(modTime.dwLowDateTime >> 8);
	timestamp[7] = static_cast<UInt8>(modTime.dwLowDateTime);

	lua_pushlstring(L, reinterpret_cast<const char *>(timestamp), 8);

	return 1;

}

int CompareFileTimestamps(lua_State *L)
{
	const UInt8 *p1 = reinterpret_cast<const UInt8 *>(lua_tostring(L, 1));
	const UInt8 *p2 = reinterpret_cast<const UInt8 *>(lua_tostring(L, 2));
	for(int i=0;i<8;i++)
	{
		if(p1[i] < p2[i])
		{
			lua_pushinteger(L, -1);
			return 1;
		}
		if(p1[i] > p2[i])
		{
			lua_pushinteger(L, 1);
			return 1;
		}
	}
	lua_pushinteger(L, 0);
	return 1;
}

int SetFileTimestamp(lua_State *L)
{
	const UInt8 *timestamp = reinterpret_cast<const UInt8 *>(lua_tostring(L, 2));

	FILETIME convertedStamp;
	convertedStamp.dwHighDateTime = (timestamp[0] << 24) | (timestamp[1] << 16) | (timestamp[2] << 8) | (timestamp[3]);
	convertedStamp.dwLowDateTime = (timestamp[4] << 24) | (timestamp[5] << 16) | (timestamp[6] << 8) | (timestamp[7]);

	HANDLE f = CreateFile(lua_tostring(L, 1), FILE_WRITE_ATTRIBUTES, FILE_SHARE_READ, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
	if(f == INVALID_HANDLE_VALUE)
		return 0;

	BOOL result = SetFileTime(f, NULL, NULL, &convertedStamp);
	CloseHandle(f);

	return 0;
}

int DestroyFile(lua_State *L)
{
	remove(lua_tostring(L, 1));
	return 0;
}

int CreateFileDirectories(lua_State *L)
{
	char path[8000];
	strcpy(path, lua_tostring(L, 1));
	size_t len = strlen(path);

	for(size_t i=0;i<len;i++)
	{
		if(path[i] == '/' || path[i] == '\\')
		{
			char oldChar = path[i];
			path[i] = '\0';
			_mkdir(path);
			path[i] = oldChar;
		}
	}

	return 0;
}

int RunProcL(lua_State *L)
{
	CHAR cmdLine[8000];
	strcpy(cmdLine, lua_tostring(L, 2));

	PROCESS_INFORMATION procInfo;
	STARTUPINFO startInfo;
	memset(&startInfo, 0, sizeof(startInfo));
	BOOL procStarted = CreateProcess(lua_tostring(L, 1), cmdLine, NULL, NULL, FALSE, CREATE_NEW_CONSOLE, NULL, NULL, &startInfo, &procInfo);

	return 0;
}


extern "C" __declspec(dllexport) int RegisterRDistro(lua_State *L)
{
	lua_pushcfunction(L, RunProcL);
	lua_setglobal(L, "RunProc");

	lua_pushcfunction(L, GenerateContentKey);
	lua_setglobal(L, "GenerateContentKey");

	lua_pushcfunction(L, GenerateAuthKey);
	lua_setglobal(L, "GenerateAuthKey");
	
	lua_pushcfunction(L, ForceSetMetaTable);
	lua_setglobal(L, "ForceSetMetaTable");

	lua_pushcfunction(L, EncodeInteger);
	lua_setglobal(L, "EncodeInteger");

	lua_pushcfunction(L, DecodeInteger);
	lua_setglobal(L, "DecodeInteger");
	
	lua_pushcfunction(L, TimeSeconds);
	lua_setglobal(L, "TimeSeconds");

	lua_pushcfunction(L, sha256);
	lua_setglobal(L, "SHA256");

	lua_pushcfunction(L, SleepThread);
	lua_setglobal(L, "SleepThread");

	lua_pushcfunction(L, CompareFileTimestamps);
	lua_setglobal(L, "CompareFileTimestamps");
	
	lua_pushcfunction(L, CreateFileDirectories);
	lua_setglobal(L, "CreateFileDirectories");
	
	lua_pushcfunction(L, SetFileTimestamp);
	lua_setglobal(L, "SetFileTimestamp");

	lua_pushcfunction(L, FileTimestamp);
	lua_setglobal(L, "FileTimestamp");

	lua_pushcfunction(L, DestroyFile);
	lua_setglobal(L, "DestroyFile");

	lua_pushcfunction(L, Sockets_SelectSingle);
	lua_setglobal(L, "Sockets_SelectSingle");

	lua_pushcfunction(L, Sockets_EnvironmentStartup);
	lua_setglobal(L, "Sockets_EnvironmentStartup");

	lua_pushcfunction(L, Sockets_EnvironmentCleanup);
	lua_setglobal(L, "Sockets_EnvironmentCleanup");
	
	lua_pushcfunction(L, Sockets_GetAddrInfoCluster);
	lua_setglobal(L, "Sockets_GetAddrInfoCluster");
	
	lua_pushcfunction(L, Sockets_AddrInfoClusterFirst);
	lua_setglobal(L, "Sockets_AddrInfoClusterFirst");
	
	lua_pushcfunction(L, Sockets_AddrInfoNext);
	lua_setglobal(L, "Sockets_AddrInfoNext");
	
	lua_pushcfunction(L, Sockets_AddrInfoFamily);
	lua_setglobal(L, "Sockets_AddrInfoFamily");

	lua_pushcfunction(L, Sockets_FreeAddrInfoCluster);
	lua_setglobal(L, "Sockets_FreeAddrInfoCluster");
	
	lua_pushcfunction(L, Sockets_CreateSocket);
	lua_setglobal(L, "Sockets_CreateSocket");
	
	lua_pushcfunction(L, Sockets_CloseSocket);
	lua_setglobal(L, "Sockets_CloseSocket");
	
	lua_pushcfunction(L, Sockets_Bind);
	lua_setglobal(L, "Sockets_Bind");
	
	lua_pushcfunction(L, Sockets_Listen);
	lua_setglobal(L, "Sockets_Listen");
	
	lua_pushcfunction(L, Sockets_Accept);
	lua_setglobal(L, "Sockets_Accept");
	
	lua_pushcfunction(L, Sockets_Send);
	lua_setglobal(L, "Sockets_Send");
	
	lua_pushcfunction(L, Sockets_TryReceive);
	lua_setglobal(L, "Sockets_TryReceive");
	
	lua_pushcfunction(L, Sockets_Shutdown);
	lua_setglobal(L, "Sockets_Shutdown");

	return 0;
}
